#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>

/**
 * 线性插值函数（单通道/三通道通用）
 * @param v1: 起始值（单通道为 uint8_t，三通道为 uint8_t[3]）
 * @param v2: 结束值
 * @param t: 插值权重 (0 ≤ t ≤ 1)
 * @param channels: 通道数（1=灰度，3=RGB）
 * @param result: 输出插值结果（长度为 channels）
 */
void linear_interp(const uint8_t *v1, const uint8_t *v2, float t, int channels, uint8_t *result) {
    for (int c = 0; c < channels; c++) {
        float val = (1 - t) * v1[c] + t * v2[c];
        // 限制范围在 0-255，避免溢出
        result[c] = (uint8_t)(val < 0 ? 0 : (val > 255 ? 255 : val));
    }
}

/**
 * 图像 4 倍线性插值放大（双线性插值）
 * @param src_data: 源图像数据（RGB 按行存储：r0,g0,b0,r1,g1,b1...）
 * @param src_w: 源图像宽度
 * @param src_h: 源图像高度
 * @param channels: 通道数（1=灰度，3=RGB）
 * @param dst_data: 输出目标图像数据（需提前分配内存：src_w*4 * src_h*4 * channels）
 */
void upscale_linear_4x(const uint8_t *src_data, int src_w, int src_h, int channels, uint8_t *dst_data) {
    int dst_w = src_w * 4;
    int dst_h = src_h * 4;

    // 遍历源图像的 2x2 像素块（最后一行/列单独处理）
    for (int y = 0; y < src_h - 1; y++) {
        for (int x = 0; x < src_w - 1; x++) {
            // 获取 2x2 源像素的指针（ul: 左上, ur: 右上, dl: 左下, dr: 右下）
            const uint8_t *ul = src_data + (y * src_w + x) * channels;
            const uint8_t *ur = src_data + (y * src_w + x + 1) * channels;
            const uint8_t *dl = src_data + ((y + 1) * src_w + x) * channels;
            const uint8_t *dr = src_data + ((y + 1) * src_w + x + 1) * channels;

            // 目标 4x4 块的左上角坐标
            int ox = x * 4;
            int oy = y * 4;

            // 缓存插值结果的临时数组
            uint8_t top[3], bottom[3], pixel[3];

            // 填充 4x4 块的每个像素（i,j: 0-4，包含边界）
            for (int i = 0; i < 5; i++) {  // 水平方向（0-4）
                for (int j = 0; j < 5; j++) {  // 垂直方向（0-4）
                    int dx = ox + i;
                    int dy = oy + j;
                    if (dx >= dst_w || dy >= dst_h) continue;  // 边界保护

                    // 跳过四个角（已直接复制，此处实际会被插值覆盖，可保留逻辑对齐Python）
                    if ((i == 0 && j == 0) || (i == 4 && j == 0) ||
                        (i == 0 && j == 4) || (i == 4 && j == 4)) {
                        // 直接复制角点
                        if (i == 0 && j == 0) memcpy(dst_data + (dy * dst_w + dx) * channels, ul, channels);
                        else if (i == 4 && j == 0) memcpy(dst_data + (dy * dst_w + dx) * channels, ur, channels);
                        else if (i == 0 && j == 4) memcpy(dst_data + (dy * dst_w + dx) * channels, dl, channels);
                        else if (i == 4 && j == 4) memcpy(dst_data + (dy * dst_w + dx) * channels, dr, channels);
                        continue;
                    }

                    // 计算插值权重
                    float tx = (float)i / 4.0f;  // 水平权重（0-1）
                    float ty = (float)j / 4.0f;  // 垂直权重（0-1）

                    // 水平插值：上边（ul→ur）、下边（dl→dr）
                    linear_interp(ul, ur, tx, channels, top);
                    linear_interp(dl, dr, tx, channels, bottom);

                    // 垂直插值：top→bottom
                    linear_interp(top, bottom, ty, channels, pixel);

                    // 写入目标像素
                    memcpy(dst_data + (dy * dst_w + dx) * channels, pixel, channels);
                }
            }
        }
    }

    // ------------------- 边界处理 -------------------
    // 处理最后一列（无右侧像素，直接复制源像素到 4 列）
    for (int y = 0; y < src_h; y++) {
        const uint8_t *src_pixel = src_data + (y * src_w + src_w - 1) * channels;
        int oy = y * 4;
        for (int j = 0; j < 4; j++) {  // 垂直方向 4 个像素
            int dy = oy + j;
            if (dy >= dst_h) break;
            for (int i = 0; i < 4; i++) {  // 水平方向 4 个像素
                int dx = dst_w - 4 + i;
                if (dx >= dst_w) break;
                memcpy(dst_data + (dy * dst_w + dx) * channels, src_pixel, channels);
            }
        }
    }

    // 处理最后一行（无下方像素，直接复制源像素到 4 行）
    for (int x = 0; x < src_w; x++) {
        const uint8_t *src_pixel = src_data + ((src_h - 1) * src_w + x) * channels;
        int ox = x * 4;
        for (int i = 0; i < 4; i++) {  // 水平方向 4 个像素
            int dx = ox + i;
            if (dx >= dst_w) break;
            for (int j = 0; j < 4; j++) {  // 垂直方向 4 个像素
                int dy = dst_h - 4 + j;
                if (dy >= dst_h) break;
                memcpy(dst_data + (dy * dst_w + dx) * channels, src_pixel, channels);
            }
        }
    }

    // 右下角 4x4 块补齐（复制源图像最后一个像素）
    const uint8_t *last_pixel = src_data + ((src_h - 1) * src_w + src_w - 1) * channels;
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            int dx = dst_w - 4 + i;
            int dy = dst_h - 4 + j;
            if (dx >= dst_w || dy >= dst_h) break;
            memcpy(dst_data + (dy * dst_w + dx) * channels, last_pixel, channels);
        }
    }
}

int main() {
    // ------------------- 配置参数 -------------------
    const char *input_path = "lena_128.png";
    const char *output_path = "lena_512_linear.png";
    // ------------------------------------------------

    // 读取源图像
    int src_w, src_h, channels;
    uint8_t *src_data = stbi_load(input_path, &src_w, &src_h, &channels, 0);
    if (!src_data) {
        fprintf(stderr, "Error: 无法读取图像 %s\n", input_path);
        return -1;
    }
    printf("源图像信息：%dx%d, 通道数：%d\n", src_w, src_h, channels);

    // 检查通道数（仅支持 1/3 通道）
    if (channels != 1 && channels != 3) {
        fprintf(stderr, "Error: 仅支持灰度(1)或RGB(3)图像\n");
        stbi_image_free(src_data);
        return -1;
    }

    // 计算目标图像尺寸并分配内存
    int dst_w = src_w * 4;
    int dst_h = src_h * 4;
    size_t dst_size = (size_t)dst_w * dst_h * channels;
    uint8_t *dst_data = (uint8_t *)malloc(dst_size);
    if (!dst_data) {
        fprintf(stderr, "Error: 内存分配失败\n");
        stbi_image_free(src_data);
        return -1;
    }

    // 执行 4 倍放大
    upscale_linear_4x(src_data, src_w, src_h, channels, dst_data);

    // 保存目标图像（PNG格式）
    int success = stbi_write_png(output_path, dst_w, dst_h, channels, dst_data, dst_w * channels);
    if (success) {
        printf("成功：图像已放大 4 倍并保存为 %s\n", output_path);
    } else {
        fprintf(stderr, "Error: 保存图像失败\n");
    }

    // 释放内存
    stbi_image_free(src_data);
    free(dst_data);

    return 0;
}